# -*- coding: utf-8 -*-
"""
WebRTC服务类
负责与前端建立WebRTC连接，传输设备数据和处理用户输入
"""

import asyncio
from aiortc import RTCPeerConnection, RTCSessionDescription, RTCIceCandidate
from aiortc import RTCDataChannel, VideoStreamTrack
from aiortc.contrib.media import MediaPlayer
import threading
from typing import Dict
from dev_server.device_context_mgr import DeviceContextMgr
from dev_server.zmq_server import ZmqServer
from dev_server.webrtc_connection import WebRTCConnection
from common.utils import LogUtil
from common.webrtc_errno import WebRTCErrno

logger = LogUtil.get_logger(__name__)

class WebRTCService:
    """WebRTC服务"""
    
    def __init__(self, zmq_mgr: ZmqServer, device_context_mgr: DeviceContextMgr):
        self.connections: Dict[str, WebRTCConnection] = {}
        self.device_context_mgr = device_context_mgr
        self.zmq_mgr = zmq_mgr
        self.running = False
        self.service_thread = None
        self.loop = None
        self.message_task = None

        self.service_thread = threading.Thread(target=self._run_loop, daemon=True)

    def __del__(self):
        self.destroy()

    def destroy(self):
        """销毁WebRTC服务"""
        self.running = False
        if self.loop and self.loop.is_running():
            asyncio.run_coroutine_threadsafe(self._async_stop(), self.loop)
        if self.service_thread:
            self.service_thread.join()
        self.message_task = None
        self.loop = None
        self.service_thread = None
        self.zmq_mgr = None
        self.device_context_mgr = None
        for connection in self.connections.values():
            connection.destroy()
        self.connections.clear()
        logger.info("WebRTC服务已销毁")
    
    async def _async_stop(self):
        """异步停止WebRTC服务"""
        # 取消消息处理任务
        if self.message_task:
            self.message_task.cancel()
            try:
                await self.message_task
            except asyncio.CancelledError:
                pass
        
        # 关闭所有连接
        for connection in self.connections.values():
            await connection.close()
        self.connections.clear()

    def start(self):
        """启动WebRTC服务"""
        # 启动事件循环线程
        self.running = True
        self.service_thread.start()
        logger.info("WebRTC服务启动成功")
    
    def is_alive(self):
        """检查服务是否运行"""
        return self.service_thread and self.service_thread.is_alive()

    def _run_loop(self):
        """在新线程中运行事件循环"""
        # 创建事件循环
        self.loop = asyncio.new_event_loop()
        asyncio.set_event_loop(self.loop)
        try:
            # 设置异步事件循环主任务
            self.loop.run_until_complete(self._async_start())
        except Exception as e:
            logger.error(f"WebRTC服务事件循环异常: {e}")
    
    async def _async_start(self):
        """异步启动WebRTC服务"""
        try:            
            # 启动消息处理循环
            self.message_task = asyncio.create_task(self.__process_mesg())
            
            # 等待任务完成
            await self.message_task
            
        except Exception as e:
            logger.error(f"异步启动WebRTC服务失败: {e}")
            raise
    
    async def __process_mesg(self):
        """消息处理循环"""
        while self.running:
            message, success = await self.zmq_mgr.async_recv_json()
            if success is False:
                continue

            msg_type = 'webrtc_error'
            msg_data = {
                'source': 'webrtc_none',
                'errnum': WebRTCErrno.EPEER.value,
                'errmsg': 'Failed to parse webrtc message',
            }
            if message is not None:
                msg_type = message.pop('type')
                msg_data = message.pop('data')

            logger.info(f"处理信令消息: type={msg_type}.")

            try:
                match msg_type:
                    case 'webrtc_offer':
                        response = await self._handle_webrtc_offer(msg_data)
                    case 'webrtc_candidate':
                        response = await self._handle_webrtc_candidate(msg_data)
                    case 'webrtc_disconnect':
                        response = await self._handle_webrtc_disconnect(msg_data)
                    case 'webrtc_error':
                        response = {
                            'type': msg_type,
                            'data': msg_data,
                        }
            except Exception as e:
                logger.error(f"处理信令消息失败: {e}")
                response = { 
                    'type': 'webrtc_error',
                    'data': {
                        'source': msg_type,
                        'errnum': WebRTCErrno.EPEER.value,
                        'errmsg': str(e),
                    }
                }
            
            logger.info(f"发送信令消息: type={response.get('type')}.")

            await self.zmq_mgr.async_send_json(response)
    
    async def _handle_webrtc_offer(self, data):
        """处理请求offer"""
        client_id = data['client_id']
        device_sn = data['device_sn']
        offer = data['offer']
        # 1. 处理offer格式
        if not isinstance(offer, dict):
            logger.error(f"failed to handle webrtc_offer: invalid offer format")
            return {
                'type': 'webrtc_error',
                'data': {
                    'source': 'webrtc_offer',
                    'errnum': WebRTCErrno.EINVAL.value,
                    'errmsg': 'Invalid offer format',
                }
            }
        
        try:
            # 2. 创建新的PeerConnection
            device_ctx = self.device_context_mgr.get_device_context(device_sn)
            if device_ctx is None:
                logger.error(f"failed to handle webrtc_offer: device {device_sn} not connected")
                return {
                    'type': 'webrtc_error',
                    'data': {
                        'source': 'webrtc_offer',
                        'errnum': WebRTCErrno.EPEER.value,
                        'errmsg': f'Device {device_sn} not connected',
                    }
                }
            webrtc_conn = WebRTCConnection(client_id, device_sn, device_ctx)
            self.connections[client_id] = webrtc_conn
            
            # 3 设置事件处理器
            await webrtc_conn.setup_pc_handlers()
            
            # 4. 设置远程描述（必须在添加轨道之前）
            if await webrtc_conn.set_remote_description(offer) is False:
                logger.error(f"failed to handle webrtc_offer: failed to set remote description")
                return {
                    'type': 'webrtc_error',
                    'data': {
                        'source': 'webrtc_offer',
                        'errnum': WebRTCErrno.EPEER.value,
                        'errmsg': 'Failed to set remote description',
                    }
                }
            
            # 5. 设置视频轨道
            if webrtc_conn.set_video_track() is False:
                logger.error(f"failed to handle webrtc_offer: failed to set video track")
                return {
                    'type': 'webrtc_error',
                    'data': {
                        'source': 'webrtc_offer',
                        'errnum': WebRTCErrno.EPEER.value,
                        'errmsg': 'Failed to set video track',
                    }
                }
                
            # 6. 创建并返回answer
            answer = await webrtc_conn.create_answer()
            if answer is None:
                logger.error(f"failed to handle webrtc_offer: failed to create answer")
                return {
                    'type': 'webrtc_error',
                    'data': {
                        'source': 'webrtc_offer',
                        'errnum': WebRTCErrno.EPEER.value,
                        'errmsg': 'Failed to create answer',
                    }
                }
            
            return {
                'type': 'webrtc_answer',
                'data': {
                    'signal_syn_seq': data['signal_ack_seq'],
                    'signal_ack_seq': data['signal_syn_seq'] + 1,
                    'client_id': client_id,
                    'device_sn': device_sn,
                    'answer': answer,
                }
            }
            
        except Exception as e:
            logger.error(f"处理webrtc_offer失败: {e}")
            return {
                'type': 'webrtc_error',
                'data': {
                    'source': 'webrtc_offer',
                    'errnum': WebRTCErrno.EPEER.value,
                    'errmsg': str(e),
                }
            }

    async def _handle_webrtc_candidate(self, data):
        """处理candidate"""
        client_id = data['client_id']
        device_sn = data['device_sn']
        candidate = data['candidate']
        add_candidate_success = False
        # 1. 检查连接是否存在
        webrtc_conn = self.connections.get(client_id, None)
        if webrtc_conn is None:
            logger.error(f"failed to handle webrtc_candidate: connection {client_id} not found")
            return {
                'type': 'webrtc_error',
                'data': {
                    'source': 'webrtc_candidate',
                    'errnum': WebRTCErrno.EPEER.value,
                    'errmsg': f'Connection {client_id} not found',
                }
            }
        if candidate is not None:
            # 2. 处理candidate格式
            if not isinstance(candidate, dict):
                logger.error(f"failed to handle webrtc_candidate: invalid candidate format")
                return {
                    'type': 'webrtc_error',
                    'data': {
                        'source': 'webrtc_candidate',
                        'errnum': WebRTCErrno.EINVAL.value,
                        'errmsg': 'Invalid candidate format',
                    }
                }
            # 3. 添加ICE候选
            if await webrtc_conn.add_ice_candidate(candidate) is False:
                logger.error(f"failed to handle webrtc_candidate: failed to add ICE candidate")
                return {
                    'type': 'webrtc_error',
                    'data': {
                        'source': 'webrtc_candidate',
                        'errnum': WebRTCErrno.EPEER.value,
                        'errmsg': 'Failed to add ICE candidate',
                    }
                }
            add_candidate_success = True
        # 4. 发送candidate
        webrtc_candidate = {
            'type': 'webrtc_candidate',
            'data': {
                'signal_syn_seq': data['signal_ack_seq'],
                'signal_ack_seq': data['signal_syn_seq'] + 1,
                'client_id': client_id,
                'device_sn': device_sn,
                'candidate': None,
                'syn_seq': data['ack_seq'],
                'ack_seq': data['syn_seq'] + 1 if add_candidate_success else data['syn_seq'],
            }
        }
        local_candidate = webrtc_conn.get_ice_candidate(webrtc_candidate['data']['syn_seq'])
        if local_candidate is not None:
            webrtc_candidate['data']['candidate'] = local_candidate
        return webrtc_candidate
    
    async def _handle_webrtc_disconnect(self, data):
        """处理断开连接"""
        client_id = data.get('client_id', None)
        device_sn = data.get('device_sn', None)
        webrtc_conn = self.connections.pop(client_id, None)
        if webrtc_conn is not None:
            await webrtc_conn.close()
        return {
            'type': 'webrtc_disconnect',
            'data': {
                'signal_syn_seq': data['signal_ack_seq'],
                'signal_ack_seq': data['signal_syn_seq'] + 1,
                'client_id': client_id,
                'device_sn': device_sn,
            }
        }
